//
//  AVLTree.swift
//  others
//
//  Created by lhc on 2024/8/26.
//

import Foundation
/// AVL树
class AVLTree<T: Comparable, E>: BalanceBinaryTree<T, E> {
    
    func add(key: T, value: E? = nil) {
        self.root = self.add(key: key, value: value, node: self.root)
    }
    
    private func add(key: T, value: E?, node: Node?) -> Node? {
        guard let node = node else {
            // 当前node为空，那么在此创建新node
            self.size += 1
            return Node(key: key, value: value)
        }
        
        
        if key < node.key {
            node.left = self.add(key: key, value: value, node: node.left)
        }else if key > node.key {
            node.right = self.add(key: key, value: value, node: node.right)
        }else {
            // key == node.key
            // 重制node的value
            node.value = value
        }
        
        // 更新node的高度值
        // 高度值是左右子树的高度最大值 + 1
        node.height = max(self.getHeightWithNode(node.left), self.getHeightWithNode(node.right)) + 1
        // 更新当前node的平衡因子
        let balanceFactor = self.getBalanceFactor(node: node)
        
        var resultNode: Node? = node
        // 判断分叉的方向
        // LL
        if balanceFactor > 1 && self.getBalanceFactor(node: resultNode?.left) >= 0 {
            resultNode = self.rightRotate(node: resultNode)
        // LR
        }else if balanceFactor > 1 && self.getBalanceFactor(node: resultNode?.left) < 0 {
            resultNode?.left = self.leftRotate(node: resultNode?.left)
            resultNode = self.rightRotate(node: resultNode)
        // RR
        }else if balanceFactor < -1 && self.getBalanceFactor(node: resultNode?.right) <= 0 {
            resultNode = self.leftRotate(node: resultNode)
        // RL
        }else if balanceFactor < -1 && self.getBalanceFactor(node: resultNode?.right) > 0 {
            resultNode?.right = self.rightRotate(node: resultNode?.right)
            resultNode = self.leftRotate(node: resultNode)
        }
        
        return resultNode
    }
    
    /// 左旋
    private func leftRotate(node: Node?) -> Node? {
        guard let node = node else {return nil}
        
        let x = node.right
        node.right = x?.left
        x?.left = node
        
        // 先更新node高度 再更新x高度 (因为左旋后node 是 x的左子树)
        node.height = max(self.getHeightWithNode(node.left), self.getHeightWithNode(node.right)) + 1
        x?.height = max(self.getHeightWithNode(x?.left), self.getHeightWithNode(x?.right)) + 1
        
        return x
    }
    
    /// 右旋
    private func rightRotate(node: Node?) -> Node? {
        guard let node = node else {return nil}
        let x = node.left
        node.left = x?.right
        x?.right = node
        
        // 先更新node高度 再更新x高度 (因为左旋后node 是 x的右子树)
        node.height = max(self.getHeightWithNode(node.left), self.getHeightWithNode(node.right)) + 1
        x?.height = max(self.getHeightWithNode(x?.left), self.getHeightWithNode(x?.right)) + 1
        
        return x
    }
    
    func removeKey(_ key: T) {
        guard !self.isEmpty() else {return}
        
        self.root = self.removeKey(key, node: self.root)
    }
    
    private func removeKey(_ key: T, node: Node?) -> Node? {
        guard let node = node else { return nil }
        
        var resultNode: Node? = nil
        
        if key < node.key {
            node.left = self.removeKey(key, node: node.left)
            resultNode = node
        }else if key > node.key {
            node.right = self.removeKey(key, node: node.right)
            resultNode = node
        }else {
            // key == node.key
            
            if node.left == nil {
                resultNode = node.right
                node.right = nil
                self.size -= 1
                
            }else if node.right == nil {
                resultNode = node.left
                node.left = nil
                self.size -= 1
            }else {
                // node.left node.right 都不为空
                // 找到node.right为根的最小node, 也就是当前node值大一点的下一个node
                let successor = self.miniNode(node: node.right!)
                successor.right = self.removeKey(successor.key, node: node.right)
                successor.left = node.left
                
                node.left = nil
                node.right = nil
                
                resultNode = successor
            }
        }
        
        if resultNode == nil {
            return nil
        }
        
        // 做平衡判断
        resultNode?.height = max(self.getHeightWithNode(resultNode?.left), self.getHeightWithNode(resultNode?.right)) + 1
        
        // 更新当前node的平衡因子
        let balanceFactor = self.getBalanceFactor(node: node)
        if balanceFactor > 1 && self.getBalanceFactor(node: resultNode?.left) >= 0 {
            resultNode = self.rightRotate(node: resultNode)
        // LR
        }else if balanceFactor > 1 && self.getBalanceFactor(node: resultNode?.left) < 0 {
            resultNode?.left = self.leftRotate(node: resultNode?.left)
            resultNode = self.rightRotate(node: resultNode)
        // RR
        }else if balanceFactor < -1 && self.getBalanceFactor(node: resultNode?.right) <= 0 {
            resultNode = self.leftRotate(node: resultNode)
        // RL
        }else if balanceFactor < -1 && self.getBalanceFactor(node: resultNode?.right) > 0 {
            resultNode?.right = self.rightRotate(node: resultNode?.right)
            resultNode = self.leftRotate(node: resultNode)
        }
        
        
        return resultNode
    }
}



